#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
上市公司营收预测模型实战演练
根据第五章内容完成实战演练部分任务
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest, f_regression
import warnings
warnings.filterwarnings('ignore')

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 数据加载与预处理
def load_and_preprocess_data():
    """加载并预处理数据"""
    print("1. 数据加载与预处理")
    
    # 加载数据
    data = pd.read_excel('上市公司营收数据.xlsx')
    print(f"数据形状: {data.shape}")
    print("数据前5行:")
    print(data.head())
    
    # 检查缺失值
    print("\n缺失值检查:")
    print(data.isnull().sum())
    
    # 处理缺失值
    if data.isnull().sum().sum() > 0:
        data = data.dropna()  # 简单起见，直接删除缺失值
        print(f"删除缺失值后数据形状: {data.shape}")
    
    # 检查数据类型
    print("\n数据类型:")
    print(data.dtypes)
    
    # 基本统计信息
    print("\n基本统计信息:")
    print(data.describe())
    
    return data

# 2. 探索性数据分析
def exploratory_data_analysis(data):
    """探索性数据分析"""
    print("\n2. 探索性数据分析")
    
    # 获取数值列
    numeric_cols = data.select_dtypes(include=[np.number]).columns.tolist()
    
    # 相关性分析
    print("\n相关性分析:")
    corr_matrix = data[numeric_cols].corr()
    print(corr_matrix)
    
    # 绘制热力图
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f')
    plt.title('特征相关性热力图')
    plt.tight_layout()
    plt.savefig('相关性热力图.png')
    plt.close()
    
    # 绘制各特征的分布
    data[numeric_cols].hist(figsize=(15, 12), bins=20)
    plt.suptitle('特征分布直方图')
    plt.tight_layout()
    plt.savefig('特征分布直方图.png')
    plt.close()
    
    # 绘制箱线图
    plt.figure(figsize=(15, 8))
    data[numeric_cols].boxplot()
    plt.title('特征箱线图')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig('特征箱线图.png')
    plt.close()
    
    return numeric_cols

# 3. 特征工程
def feature_engineering(data, numeric_cols):
    """特征工程"""
    print("\n3. 特征工程")
    
    # 假设最后一列是目标变量（营收）
    target_col = numeric_cols[-1]
    feature_cols = numeric_cols[:-1]
    
    print(f"目标变量: {target_col}")
    print(f"特征变量: {feature_cols}")
    
    X = data[feature_cols]
    y = data[target_col]
    
    # 特征选择
    selector = SelectKBest(score_func=f_regression, k='all')
    selector.fit(X, y)
    
    # 打印特征重要性
    feature_scores = pd.DataFrame({
        'Feature': feature_cols,
        'Score': selector.scores_
    })
    feature_scores = feature_scores.sort_values(by='Score', ascending=False)
    print("\n特征重要性:")
    print(feature_scores)
    
    # 绘制特征重要性
    plt.figure(figsize=(10, 6))
    sns.barplot(x='Score', y='Feature', data=feature_scores)
    plt.title('特征重要性')
    plt.tight_layout()
    plt.savefig('特征重要性.png')
    plt.close()
    
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    X_scaled = pd.DataFrame(X_scaled, columns=feature_cols)
    
    return X_scaled, y, feature_cols, target_col

# 4. 数据集划分
def split_dataset(X, y):
    """划分数据集"""
    print("\n4. 数据集划分")
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42
    )
    
    print(f"训练集大小: {X_train.shape}")
    print(f"测试集大小: {X_test.shape}")
    
    return X_train, X_test, y_train, y_test

# 5. 模型建立与训练
def build_and_train_models(X_train, X_test, y_train, y_test):
    """建立并训练模型"""
    print("\n5. 模型建立与训练")
    
    models = {
        '线性回归': LinearRegression(),
        '岭回归': Ridge(),
        'Lasso回归': Lasso(),
        '随机森林': RandomForestRegressor(random_state=42)
    }
    
    results = {}
    
    for name, model in models.items():
        print(f"\n训练 {name} 模型...")
        
        # 训练模型
        model.fit(X_train, y_train)
        
        # 预测
        y_pred_train = model.predict(X_train)
        y_pred_test = model.predict(X_test)
        
        # 评估指标
        mse_train = mean_squared_error(y_train, y_pred_train)
        mse_test = mean_squared_error(y_test, y_pred_test)
        rmse_train = np.sqrt(mse_train)
        rmse_test = np.sqrt(mse_test)
        mae_train = mean_absolute_error(y_train, y_pred_train)
        mae_test = mean_absolute_error(y_test, y_pred_test)
        r2_train = r2_score(y_train, y_pred_train)
        r2_test = r2_score(y_test, y_pred_test)
        
        # 交叉验证
        cv_scores = cross_val_score(model, X_train, y_train, cv=5, scoring='r2')
        
        results[name] = {
            'model': model,
            'mse_train': mse_train,
            'mse_test': mse_test,
            'rmse_train': rmse_train,
            'rmse_test': rmse_test,
            'mae_train': mae_train,
            'mae_test': mae_test,
            'r2_train': r2_train,
            'r2_test': r2_test,
            'cv_mean': cv_scores.mean(),
            'cv_std': cv_scores.std()
        }
        
        print(f"{name} 训练集 R²: {r2_train:.4f}")
        print(f"{name} 测试集 R²: {r2_test:.4f}")
        print(f"{name} 测试集 RMSE: {rmse_test:.4f}")
        print(f"{name} 测试集 MAE: {mae_test:.4f}")
        print(f"{name} 交叉验证 R²: {cv_scores.mean():.4f} (±{cv_scores.std():.4f})")
    
    return results

# 6. 模型评估与比较
def evaluate_and_compare_models(results, X_test, y_test, feature_cols):
    """评估并比较模型"""
    print("\n6. 模型评估与比较")
    
    # 创建比较表格
    comparison = pd.DataFrame({
        'Model': list(results.keys()),
        'Train R²': [results[name]['r2_train'] for name in results.keys()],
        'Test R²': [results[name]['r2_test'] for name in results.keys()],
        'Test RMSE': [results[name]['rmse_test'] for name in results.keys()],
        'Test MAE': [results[name]['mae_test'] for name in results.keys()],
        'CV R²': [results[name]['cv_mean'] for name in results.keys()],
        'CV Std': [results[name]['cv_std'] for name in results.keys()]
    })
    
    print("\n模型比较:")
    print(comparison)
    
    # 绘制模型比较图
    plt.figure(figsize=(15, 10))
    
    # R²比较
    plt.subplot(2, 2, 1)
    sns.barplot(x='Model', y='Test R²', data=comparison)
    plt.title('模型测试集 R² 比较')
    plt.xticks(rotation=45)
    
    # RMSE比较
    plt.subplot(2, 2, 2)
    sns.barplot(x='Model', y='Test RMSE', data=comparison)
    plt.title('模型测试集 RMSE 比较')
    plt.xticks(rotation=45)
    
    # MAE比较
    plt.subplot(2, 2, 3)
    sns.barplot(x='Model', y='Test MAE', data=comparison)
    plt.title('模型测试集 MAE 比较')
    plt.xticks(rotation=45)
    
    # 交叉验证R²比较
    plt.subplot(2, 2, 4)
    sns.barplot(x='Model', y='CV R²', data=comparison)
    plt.title('模型交叉验证 R² 比较')
    plt.xticks(rotation=45)
    
    plt.tight_layout()
    plt.savefig('模型比较.png')
    plt.close()
    
    # 选择最佳模型
    best_model_name = comparison.loc[comparison['Test R²'].idxmax(), 'Model']
    print(f"\n最佳模型: {best_model_name}")
    
    return best_model_name, comparison

# 7. 特征重要性分析
def analyze_feature_importance(results, best_model_name, feature_cols):
    """分析特征重要性"""
    print(f"\n7. 分析 {best_model_name} 的特征重要性")
    
    best_model = results[best_model_name]['model']
    
    if best_model_name == '随机森林':
        # 随机森林的特征重要性
        importances = best_model.feature_importances_
        feature_importance = pd.DataFrame({
            'Feature': feature_cols,
            'Importance': importances
        }).sort_values(by='Importance', ascending=False)
        
        print("随机森林特征重要性:")
        print(feature_importance)
        
        # 绘制特征重要性
        plt.figure(figsize=(10, 6))
        sns.barplot(x='Importance', y='Feature', data=feature_importance)
        plt.title(f'{best_model_name} 特征重要性')
        plt.tight_layout()
        plt.savefig(f'{best_model_name}_特征重要性.png')
        plt.close()
    
    elif best_model_name in ['线性回归', '岭回归', 'Lasso回归']:
        # 线性模型的系数
        coefficients = best_model.coef_
        feature_importance = pd.DataFrame({
            'Feature': feature_cols,
            'Coefficient': coefficients
        })
        # 添加绝对值列用于排序
        feature_importance['Abs_Coefficient'] = feature_importance['Coefficient'].abs()
        feature_importance = feature_importance.sort_values(by='Abs_Coefficient', ascending=False)
        # 删除临时列
        feature_importance = feature_importance.drop('Abs_Coefficient', axis=1)
        
        print(f"{best_model_name} 特征系数:")
        print(feature_importance)
        
        # 绘制特征系数
        plt.figure(figsize=(10, 6))
        sns.barplot(x='Coefficient', y='Feature', data=feature_importance)
        plt.title(f'{best_model_name} 特征系数')
        plt.tight_layout()
        plt.savefig(f'{best_model_name}_特征系数.png')
        plt.close()
    
    return feature_importance

# 8. 模型优化
def optimize_model(X_train, y_train, best_model_name):
    """优化最佳模型"""
    print(f"\n8. 优化 {best_model_name} 模型")
    
    if best_model_name == '随机森林':
        # 随机森林参数调优
        param_grid = {
            'n_estimators': [50, 100, 200],
            'max_depth': [None, 10, 20, 30],
            'min_samples_split': [2, 5, 10],
            'min_samples_leaf': [1, 2, 4]
        }
        model = RandomForestRegressor(random_state=42)
    
    elif best_model_name == '岭回归':
        # 岭回归参数调优
        param_grid = {
            'alpha': [0.01, 0.1, 1.0, 10.0, 100.0]
        }
        model = Ridge()
    
    elif best_model_name == 'Lasso回归':
        # Lasso回归参数调优
        param_grid = {
            'alpha': [0.01, 0.1, 1.0, 10.0, 100.0]
        }
        model = Lasso()
    
    else:  # 线性回归
        print("线性回归无需参数调优")
        return None
    
    # 网格搜索
    grid_search = GridSearchCV(
        model, param_grid, cv=5, 
        scoring='r2', n_jobs=-1, verbose=1
    )
    grid_search.fit(X_train, y_train)
    
    print(f"最佳参数: {grid_search.best_params_}")
    print(f"最佳 R²: {grid_search.best_score_:.4f}")
    
    return grid_search.best_estimator_

# 9. 最终模型评估
def final_model_evaluation(best_model, X_test, y_test):
    """最终模型评估"""
    print("\n9. 最终模型评估")
    
    # 预测
    y_pred = best_model.predict(X_test)
    
    # 评估指标
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    
    print(f"最终模型测试集 R²: {r2:.4f}")
    print(f"最终模型测试集 RMSE: {rmse:.4f}")
    print(f"最终模型测试集 MAE: {mae:.4f}")
    
    # 绘制预测值与实际值对比图
    plt.figure(figsize=(10, 6))
    plt.scatter(y_test, y_pred, alpha=0.5)
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--')
    plt.xlabel('实际值')
    plt.ylabel('预测值')
    plt.title('预测值与实际值对比')
    plt.tight_layout()
    plt.savefig('预测值与实际值对比.png')
    plt.close()
    
    # 绘制残差图
    residuals = y_test - y_pred
    plt.figure(figsize=(10, 6))
    plt.scatter(y_pred, residuals, alpha=0.5)
    plt.axhline(y=0, color='r', linestyle='--')
    plt.xlabel('预测值')
    plt.ylabel('残差')
    plt.title('残差图')
    plt.tight_layout()
    plt.savefig('残差图.png')
    plt.close()
    
    return {
        'mse': mse,
        'rmse': rmse,
        'mae': mae,
        'r2': r2
    }

# 10. 生成报告
def generate_report(data, results, best_model_name, feature_importance, final_metrics):
    """生成实战报告"""
    print("\n10. 生成实战报告")
    
    report = f"""
# 上市公司营收预测模型实战报告

## 一、案例介绍

在金融领域，上市公司营收预测非常重要，因为它直接反映了公司的运营情况和未来发展趋势。投资者和金融数据分析师通常关注上市公司营收预测，以评估公司业绩表现和投资潜力。

## 二、实战目标

建立上市公司营收预测模型，评价模型的性能，探索影响上市公司营收的关键因素。

## 三、实战任务

### 1. 设计指标，收集与整理数据

本次实战使用的数据集包含{data.shape[0]}个样本和{data.shape[1]}个特征。数据集包含以下特征：

{', '.join(data.columns.tolist())}

数据预处理步骤：
- 检查并处理缺失值
- 检查数据类型
- 基本统计分析

### 2. 建立不少于两种模型来预测上市公司的营收，评价模型性能

本次实战建立了以下四种模型：
- 线性回归
- 岭回归
- Lasso回归
- 随机森林

模型性能比较：

| 模型 | 训练集 R² | 测试集 R² | 测试集 RMSE | 测试集 MAE | 交叉验证 R² |
|------|----------|----------|------------|------------|------------|
"""
    
    for name in results.keys():
        report += f"| {name} | {results[name]['r2_train']:.4f} | {results[name]['r2_test']:.4f} | {results[name]['rmse_test']:.4f} | {results[name]['mae_test']:.4f} | {results[name]['cv_mean']:.4f} |\n"
    
    report += f"""
根据模型性能比较，选择{best_model_name}作为最佳模型。

### 3. 确定最终的预测模型，探索影响上市公司营收的关键因素

最终模型：{best_model_name}

最终模型性能：
- R²: {final_metrics['r2']:.4f}
- RMSE: {final_metrics['rmse']:.4f}
- MAE: {final_metrics['mae']:.4f}

影响上市公司营收的关键因素：

| 特征 | 重要性/系数 |
|------|------------|
"""
    
    for _, row in feature_importance.iterrows():
        report += f"| {row['Feature']} | {row['Importance' if 'Importance' in row else 'Coefficient']:.4f} |\n"
    
    report += """
### 4. 结论与建议

通过本次实战演练，我们成功建立了上市公司营收预测模型，并确定了影响营收的关键因素。主要结论如下：

1. 模型性能：{best_model_name}模型表现最佳，能够较好地预测上市公司营收。
2. 关键因素：从特征重要性分析可以看出，[根据实际结果填写]是影响上市公司营收的最重要因素。
3. 建议：投资者和分析师可以重点关注这些关键因素，以更好地评估公司的营收表现和投资潜力。

## 四、附录

### 模型可视化结果

1. 相关性热力图：展示了各特征之间的相关性。
2. 特征分布直方图：展示了各特征的分布情况。
3. 特征箱线图：展示了各特征的分布和异常值情况。
4. 特征重要性图：展示了各特征对营收预测的重要性。
5. 模型比较图：展示了不同模型的性能比较。
6. 预测值与实际值对比图：展示了模型预测的准确性。
7. 残差图：展示了模型预测的残差分布。

这些可视化结果帮助我们更好地理解数据和模型性能。
"""
    
    # 保存报告
    with open('上市公司营收预测模型实战报告.md', 'w', encoding='utf-8') as f:
        f.write(report)
    
    print("实战报告已生成: 上市公司营收预测模型实战报告.md")
    
    return report

# 主函数
def main():
    """主函数"""
    print("上市公司营收预测模型实战演练")
    print("=" * 50)
    
    # 1. 数据加载与预处理
    data = load_and_preprocess_data()
    
    # 2. 探索性数据分析
    numeric_cols = exploratory_data_analysis(data)
    
    # 3. 特征工程
    X, y, feature_cols, target_col = feature_engineering(data, numeric_cols)
    
    # 4. 数据集划分
    X_train, X_test, y_train, y_test = split_dataset(X, y)
    
    # 5. 模型建立与训练
    results = build_and_train_models(X_train, X_test, y_train, y_test)
    
    # 6. 模型评估与比较
    best_model_name, comparison = evaluate_and_compare_models(results, X_test, y_test, feature_cols)
    
    # 7. 特征重要性分析
    feature_importance = analyze_feature_importance(results, best_model_name, feature_cols)
    
    # 8. 模型优化
    optimized_model = optimize_model(X_train, y_train, best_model_name)
    
    # 9. 最终模型评估
    if optimized_model is not None:
        final_metrics = final_model_evaluation(optimized_model, X_test, y_test)
    else:
        final_metrics = final_model_evaluation(results[best_model_name]['model'], X_test, y_test)
    
    # 10. 生成报告
    report = generate_report(data, results, best_model_name, feature_importance, final_metrics)
    
    print("\n实战演练完成！")
    print("=" * 50)

if __name__ == "__main__":
    main()